package com.soup.faker.wrapper;

import com.soup.faker.api.Generator;
import com.soup.faker.api.Interceptor;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.IntSupplier;

public class TreeGenerator implements Generator<List<Map<String, Object>>> {

    private final Generator<Map<String, Object>> elementGenerator;
    private final String key;
    private final String parentKey;
    private Object root;

    private IntSupplier depthSupplier;
    private IntSupplier sizeSupplier;

    public TreeGenerator(Generator<Map<String, Object>> elementGenerator, String key, String parentKey) {
        this.elementGenerator = elementGenerator;
        this.key = key;
        this.parentKey = parentKey;
    }

    public TreeGenerator(Generator<Map<String, Object>> elementGenerator, String key, String parentKey, Object root, IntSupplier depthSupplier, IntSupplier sizeSupplier) {
        this.elementGenerator = elementGenerator;
        this.key = key;
        this.parentKey = parentKey;
        this.root = root;
        this.depthSupplier = depthSupplier;
        this.sizeSupplier = sizeSupplier;
    }

    @Override
    public List<Map<String, Object>> get() {
        int depth = depthSupplier.getAsInt();
        ListGenerator<Map<String, Object>> listGenerator = new ListGenerator<>(elementGenerator, sizeSupplier);
        Map<String, Object> rootMap = elementGenerator.get();
        if(root != null) {
            rootMap.put(key, root);
        }
        List<Map<String, Object>> result = new ArrayList<>();
        result.add(rootMap);
        result.addAll(generateChildren(listGenerator, root, depth));
        return result;
    }

    private List<Map<String, Object>> generateChildren(ListGenerator<Map<String, Object>> listGenerator, Object parentId, int depth) {
        Interceptor<List<Map<String, Object>>> interceptor = new Interceptor<List<Map<String, Object>>>() {
            @Override
            public void afterReturn(ValueWrapper<List<Map<String, Object>>> valueWrapper, Generator<List<Map<String, Object>>> generator) {
                valueWrapper.ifPresent(list -> list.forEach(map -> map.put(parentKey, parentId)));
            }
        };
        List<Map<String, Object>> result = new ArrayList<>();
        SandwichGenerator<List<Map<String, Object>>> wrapper = new SandwichGenerator<>(listGenerator, Collections.singletonList(interceptor));
        List<Map<String, Object>> parentList = wrapper.get();
        if(parentList != null && depth > 0) {
            result.addAll(parentList);
            for (Map<String, Object> parent : parentList) {
                result.addAll(generateChildren(listGenerator, parent.get(key), depth));
            }
        }
        return result;
    }

}
